{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 注意力机制",
   "id": "a256d2d8a29856de"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-26T00:30:33.403053Z",
     "start_time": "2025-05-26T00:30:33.395245Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "\n",
    "# i表示输入的一个seq，i1、i2、i3、i4表示seq中的4个token，每个token的词向量是3维向量\n",
    "i1 = torch.rand(1, 3)\n",
    "i2 = torch.rand(1, 3)\n",
    "i3 = torch.rand(1, 3)\n",
    "i4 = torch.rand(1, 3)\n",
    "\n",
    "# 合并i1、i2、i3、i4\n",
    "i = torch.cat([i1, i2, i3, i4], dim=0) # i的形状是(4, 3)，表示i中有4个token，每个token的形状是的向量长度为3\n",
    "\n",
    "# w是二维矩阵\n",
    "wq = torch.rand(3, 2)\n",
    "wk = torch.rand(3, 2)\n",
    "wv = torch.rand(3, 2)\n",
    "\n",
    "q = torch.matmul(i, wq)  # wq的形状是(3, 2), i的形状(4, 3), q的形状是(4, 2)\n",
    "k = torch.matmul(i, wk)  # wk的形状是(3, 2), i的形状(4, 3), k的形状是(4, 2)\n",
    "v = torch.matmul(i, wv)  # wv的形状是(3, 2), i的形状(4, 3), v的形状是(4, 2)\n",
    "qk = torch.matmul(q, k.T)  # q的形状是(4, 2), k.T的形状是(2, 4), qk的形状是(4, 4)\n",
    "\n",
    "qk"
   ],
   "id": "65399856fb7a9b9c",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.9055, 3.1873, 2.6239, 1.1294],\n",
       "        [3.1858, 5.3275, 4.3848, 1.8883],\n",
       "        [2.6061, 4.3597, 3.5895, 1.5445],\n",
       "        [1.1735, 1.9558, 1.6051, 0.6965]])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 27
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-26T00:30:33.456401Z",
     "start_time": "2025-05-26T00:30:33.452486Z"
    }
   },
   "cell_type": "code",
   "source": [
    "mask = torch.tril(torch.ones(4, 4))\n",
    "mask"
   ],
   "id": "d31f0e84b66b90c4",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0., 0.],\n",
       "        [1., 1., 0., 0.],\n",
       "        [1., 1., 1., 0.],\n",
       "        [1., 1., 1., 1.]])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 28
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-26T00:30:33.518188Z",
     "start_time": "2025-05-26T00:30:33.514719Z"
    }
   },
   "cell_type": "code",
   "source": [
    "attention_score = qk.masked_fill(mask == 0, -1e9)\n",
    "attention_score"
   ],
   "id": "6054dc9a1b4eca81",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.9055e+00, -1.0000e+09, -1.0000e+09, -1.0000e+09],\n",
       "        [ 3.1858e+00,  5.3275e+00, -1.0000e+09, -1.0000e+09],\n",
       "        [ 2.6061e+00,  4.3597e+00,  3.5895e+00, -1.0000e+09],\n",
       "        [ 1.1735e+00,  1.9558e+00,  1.6051e+00,  6.9655e-01]])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 29
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-26T00:30:33.585507Z",
     "start_time": "2025-05-26T00:30:33.581854Z"
    }
   },
   "cell_type": "code",
   "source": [
    "attention_weight = torch.softmax(attention_score, dim=0)\n",
    "\n",
    "o = torch.matmul(attention_weight, v)  # qk的形状是(4, 4), v的形状是(4, 2), o的形状是(4, 2), 输出是q的行，v的列, q的行等于i的行，这就对上了\n",
    "o"
   ],
   "id": "6c950b653bd20aea",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0494, 0.0576],\n",
       "        [0.7192, 0.7591],\n",
       "        [0.9722, 0.9168],\n",
       "        [0.3143, 0.4674]])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 30
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
